#include "BasisTransform.h"

#include "DynamStateBasisWrapper.h"

#include <Math.d/FullSquareMatrix.h>

#include <cmath>
#include <algorithm>

#include "Lapack32.d/dpstrf.h"

extern "C" {
  // LAPACK: Inverse a triangular matrix
  void _FORTRAN(dtrtri)(const char * uplo, const char * diag, const int * n, double * a, const int * lda, int * info);
  
  // LAPACK: QR factorization with column pivoting (A*P = Q*R)
  void _FORTRAN(dgeqp3)(const int* m, const int* n, double* a, const int* lda, int* jptv,
                        double* tau, double* work, const int* lwork, int* info);

  // LAPACK: Generates explicitely the orthogonal Q generated by DGEQP3
  void _FORTRAN(dorgqr)(const int* m, const int* n, const int* k, double* a, const int* lda,
                        const double * tau, double* work, const int* lwork, int* info);
}

namespace Pita {

GramSchmidtOrtho::GramSchmidtOrtho(size_t vectorSize, double energyTol) :
  BasisTransform(vectorSize), energyTol_(energyTol)
{}

void
BasisCopy::inputValueIs(const DynamStateBasis & iv) {
  output()->stateBasisDel();
  output()->lastStateBasisIs(&iv);
}

void
GramSchmidtOrtho::inputValueIs(const DynamStateBasis & iv) {
  output()->stateBasisDel();

  double maxEnergy = 0.0;
  for (DynamStateBasis::IteratorConst it = iv.state(); it; ++it) {
    DynamState state = *it;
    maxEnergy = std::max(state * state, maxEnergy);
  }
  double absEnergyTol = maxEnergy * energyTol_;
  
  // Modified Gram-Schmidt procedure 
  for (DynamStateBasis::IteratorConst it = iv.state(); it; ++it) {
    DynamState state = *it;
    
    for (DynamStateBasis::IteratorConst jt = output()->state(); jt; ++jt) {
      double dot = state * (*jt);
      state.linAdd(-dot, *jt);
    }

    double finalEnergy = state * state;
    if (finalEnergy > absEnergyTol) {
      state /= std::sqrt(finalEnergy);
      output()->lastStateIs(state); 
    }
  }
}

CholeskyOrtho::CholeskyOrtho(size_t vectorSize, double energyTol) :
  BasisTransform(vectorSize),
  energyTol_(energyTol)
{}

void
CholeskyOrtho::inputValueIs(const DynamStateBasis & iv) {
  if (iv.stateCount() == 0) {
    output()->stateBasisDel();
    return;
  }

  // Assemble normal matrix: W = V^T * V
  // (Lower triangular part only == upper-triangular part only in Fortran indexing)
  FullSquareMatrix normalMatrix(iv.stateCount());
  const int arraySize = normalMatrix.dim();
  
  for (int row = 0; row < arraySize; ++row) {
    DynamState rowState = iv.state(row);
    for (int col = 0; col <= row; ++col) {
      normalMatrix[row][col] = rowState * iv.state(col);
    }
  }
 
  // Determine absolute tolerance: relativeTolerance * max(diag(W))
  double maxEnergy = 0.0;
  for (int i = 0; i < arraySize; ++i) {
    maxEnergy = std::max(normalMatrix[i][i], maxEnergy);
  }
  const double absEnergyTol = maxEnergy * energyTol_;
 
  // Setup LAPACK
  const char uplo = 'U';   // Lower triangular in C indexing == upper triangular in Fortran indexing 
  SimpleBuffer<int> permutation(arraySize);
  int factorRank;

  // Pivoted Cholesky factorization: P^T * W * P = U^T * U, with U upper-triangular
  {
    SimpleBuffer<double> workspace(2 * arraySize);
    
    int info;
    _FORTRAN(dpstrf)(&uplo, &arraySize, normalMatrix.data(), &arraySize,
                     permutation.array(), &factorRank, &absEnergyTol, workspace.array(), &info);
    assert(info >= 0);
  }

  // Inverse Cholesky factor: U^{-1}
  {
    const char diag = 'N';

    int info;
    _FORTRAN(dtrtri)(&uplo, &diag, &factorRank, normalMatrix.data(), &arraySize, &info);
    assert(info == 0);
  }

  // Compute orthogonal basis: Q = V * P * U^{-1}
  output()->stateBasisDel();
  for (int s = 0; s < factorRank; ++s) {
    DynamState newState(iv.vectorSize(), 0.0);
    for (int d = 0; d <= s; ++d) {
      int index = permutation[d] - 1; // Fortran to C indexing
      newState.linAdd(normalMatrix[s][d], iv.state(index));
    }
    output()->lastStateIs(newState);
  }
}

QROrtho::QROrtho(size_t vectorSize, double energyTol) :
  BasisTransform(vectorSize),
  energyTol_(energyTol),
  data_(),
  permutation_(),
  tau_(),
  workspace_()
{}

void
QROrtho::inputValueIs(const DynamStateBasis & iv) {
  if (iv.stateCount() == 0) {
    output()->stateBasisDel();
    return;
  }

  // Allocate memory
  const int stateSize = 2 * iv.vectorSize();
  const int stateCount = iv.stateCount();

  data_.sizeIs(stateSize * stateCount);
  permutation_.sizeIs(stateCount);
  tau_.sizeIs(stateCount);

  // Fill memory
  double * stateBufferStart = data_.array();
  for (DynamStateBasis::IteratorConst it = iv.state(); it; ++it) {
    bufferStateCopy(*it, stateBufferStart);
    stateBufferStart += stateSize;
  }
  
  std::fill_n(permutation_.array(), permutation_.size(), 0); // All columns are free

  // Perform factorization: L * P = Q * R
  const int query = -1;
  double idealSize;
  int info;
 
  // Query for ideal workspace size 
  _FORTRAN(dgeqp3)(&stateSize, &stateCount, data_.array(), &stateSize,
                   permutation_.array(), tau_.array(),
                   &idealSize, &query, &info);
  assert(info == 0);
  
  int workspaceSize = static_cast<int>(idealSize);
  workspace_.sizeIs(workspaceSize);
  
  _FORTRAN(dgeqp3)(&stateSize, &stateCount, data_.array(), &stateSize,
                   permutation_.array(), tau_.array(),
                   workspace_.array(), &workspaceSize, &info);
  assert(info == 0);
  
  // Determine numerical rank 
  int numericalRank = 1;
  {
    const int diagTermStride = stateSize + 1;
    const double diagTermTol = std::abs(data_[0]) * std::sqrt(energyTol_);
    
    const double * diagTermAddr = data_.array() + diagTermStride;
    while (numericalRank < stateCount && std::abs(*diagTermAddr) > diagTermTol) {
      ++numericalRank;
      diagTermAddr += diagTermStride;
    }
  }

  // Assemble explicitely relevant part of Q
  _FORTRAN(dorgqr)(&stateSize, &numericalRank, &numericalRank,
                   data_.array(), &stateSize, tau_.array(),
                   workspace_.array(), &workspaceSize, &info);
  assert(info == 0);

  // Fill the output with the columns of Q
  output()->stateBasisDel();
  DynamStateBasisWrapper::Ptr wrapper = DynamStateBasisWrapper::New(iv.vectorSize(), numericalRank, data_.array());
  output()->lastStateBasisIs(wrapper.ptr());
}

} /* end namespace Pita */
